{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keeping long-term memories about collaboration \n",
    "\n",
    "![Short-term Memory vs. Long-term Memory](images/short-term%20memory%20vs%20long-term%20memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Semantic Memory\n",
    "\n",
    "![Semantic Memory](images/semantic-memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unstructured Memory\n",
    "\n",
    "![Unstructured Memory](images/unstructured-memory.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define a proper User Profile schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class UserProfile(BaseModel):\n",
    "    name: str = Field(\n",
    "        description=\"The user's preferred name.\",\n",
    "    )\n",
    "    profession: str = Field(\n",
    "        description=\"The user's profession or job title.\",\n",
    "    )\n",
    "    seniority: Optional[str] = Field(\n",
    "        default=None,\n",
    "        description=\"The user's seniority level, e.g., 'mid-level', 'senior', etc. If unknown, guess it based on profession description\",\n",
    "    )\n",
    "    languages: List[str] = Field(\n",
    "        description=\"A list of programming languages the user is proficient in.\",\n",
    "    )\n",
    "    frameworks: List[str] = Field(\n",
    "        description=\"A list of frameworks or major technologies the user uses.\",\n",
    "    )\n",
    "    current_project: Optional[str] = Field(\n",
    "        default=None,\n",
    "        description=\"The user's current main project.\",\n",
    "    )\n",
    "    skills: List[str] = Field(\n",
    "        description=\"A list of the user's technical or professional skills.\",\n",
    "    )\n",
    "    current_interest: Optional[str] = Field(\n",
    "        default=None,\n",
    "        description=\"The user's current main interest or focus, e.g., a topic or technology they are actively exploring or asking about.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Structured Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "structured_model = model.with_structured_output(UserProfile)\n",
    "\n",
    "profile = structured_model.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "\"\"\")])\n",
    "print(profile.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Structure Output\n",
    "\n",
    "![Structure Output](images/structured-output.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's recreate a simple profile in memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "long_term_memory = InMemoryStore()\n",
    "\n",
    "user_id = \"1\"\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "\n",
    "profile = {\n",
    "  \"name\": \"Evgeny\"\n",
    "}\n",
    "\n",
    "long_term_memory.put(namespace, key, profile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat bot with structured profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import StateGraph, MessagesState, START, END\n",
    "from langchain_core.runnables.config import RunnableConfig\n",
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "### Nodes\n",
    "\n",
    "def chat(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "\n",
    "    # Retrieve memory from the store\n",
    "    profile = store.get((user_id, \"memory\"), \"profile\")\n",
    "\n",
    "    # Extract the actual memory content if it exists and add a prefix\n",
    "    if profile:\n",
    "        profile_content = profile.value\n",
    "    else:\n",
    "        profile_content = None\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    You are a helpful assistant with memory capabilities.\n",
    "    If user-specific memory is available, use it to personalize \n",
    "    your responses based on what you know about the user.\n",
    "    \n",
    "    Your goal is to provide relevant, friendly, and tailored \n",
    "    assistance that reflects the user’s preferences, context, and past interactions.\n",
    "\n",
    "    If the user’s name or relevant personal context is available, always personalize your responses by:\n",
    "        – Addressing the user by name (e.g., \"Sure, Bob...\") when appropriate\n",
    "        – Referencing known projects, tools, or preferences (e.g., \"your MCP  server typescript based project\")\n",
    "        – Adjusting the tone to feel friendly, natural, and directly aimed at the user\n",
    "\n",
    "    Avoid generic phrasing when personalization is possible. For example, instead of \"In TypeScript apps...\" say \"Since your project is built with TypeScript...\"\n",
    "\n",
    "    Use personalization especially in:\n",
    "        – Greetings and transitions\n",
    "        – Help or guidance tailored to tools and frameworks the user uses\n",
    "        – Follow-up messages that continue from past context\n",
    "\n",
    "    Always ensure that personalization is based only on known user details and not assumed.\n",
    "    Tailor your answers to level of seniority of the user if it's known.\n",
    "    \n",
    "    The user’s profile (which may be empty) is provided as JSON (extract all data and use it): {profile_content}\n",
    "    \"\"\"\n",
    "\n",
    "    response = model.invoke([SystemMessage(content=system_msg)] + state[\"messages\"])\n",
    "\n",
    "    return {\"messages\": response}\n",
    "\n",
    "\n",
    "def update_memory(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "    namespace = (user_id, \"memory\")\n",
    "    key = \"profile\"\n",
    "    profile = store.get(namespace, key)\n",
    "        \n",
    "    if profile:\n",
    "        profile_content = profile.value\n",
    "    else:\n",
    "        profile_content = None\n",
    "\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    Update the user profile using the user's chat history.\n",
    "    Save this for future reference. If a profile already exists, just update it.\n",
    "    Here is the current profile (it might be empty): {profile_content}\n",
    "    \"\"\"\n",
    "    \n",
    "    updated_profile = structured_model.invoke([SystemMessage(content=system_msg)] + state['messages'])\n",
    "    store.put(namespace, key, updated_profile.model_dump())\n",
    "\n",
    "\n",
    "# Define the graph\n",
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(\"chat\", chat)\n",
    "builder.add_node(\"update_memory\", update_memory)\n",
    "builder.add_edge(START, \"chat\")\n",
    "builder.add_edge(\"chat\", \"update_memory\")\n",
    "builder.add_edge(\"update_memory\", END)\n",
    "\n",
    "short_term_memory = MemorySaver()\n",
    "\n",
    "graph = builder.compile(checkpointer=short_term_memory, store=long_term_memory)\n",
    "\n",
    "display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = long_term_memory.get(namespace, key)\n",
    "memory.dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "Hey, I am a junior software engineer and I need your help with my project. \n",
    "This is a TaskManager Java Spring Boot application.\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "profile = long_term_memory.get(namespace, key)\n",
    "profile.dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "I do not get how to set up security configuration!\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "profile = long_term_memory.get(namespace, key)\n",
    "profile.dict()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
